#include <assert.h>
#include <vector>
#include <iostream>
#include <cuda_fp16.h>
#include "custom_plugins.h"
#include "cuda_utils.h"

namespace Tn {
    template<typename T>
    void write(char *&buffer, const T &val) {
        *reinterpret_cast<T *>(buffer) = val;
        buffer += sizeof(T);
    }

    template<typename T>
    void read(const char *&buffer, T &val) {
        val = *reinterpret_cast<const T *>(buffer);
        buffer += sizeof(T);
    }
}

namespace nvinfer1 {
    CustomSlicePlugin::CustomSlicePlugin(const int *start, const int *stride) {
        size_t len = sizeof(int) * 4;
        memcpy(mStart, start, len);
        memcpy(mStride, stride, len);
    }

    CustomSlicePlugin::CustomSlicePlugin(const void *data, size_t length) {
        using namespace Tn;
        const char *d = reinterpret_cast<const char *>(data), *a = d;
        size_t len = sizeof(int) * 4;
        memcpy(mStart, d, len);
        d += len;
        memcpy(mStride, d, len);
        d += len;

        assert(d == a + length);
    }

    CustomSlicePlugin::~CustomSlicePlugin() {
    }

    int CustomSlicePlugin::getNbOutputs() const {
        return 1;
    }

    void CustomSlicePlugin::serialize(void *buffer) const {
        using namespace Tn;
        char *d = static_cast<char *>(buffer), *a = d;
        size_t len = sizeof(int) * 4;
        memcpy(d, mStart, len);
        d += len;
        memcpy(d, mStride, len);
        d += len;
        assert(d == a + getSerializationSize());
    }

    size_t CustomSlicePlugin::getSerializationSize() const {
        return sizeof(int) * 4 * 2;
    }

    int CustomSlicePlugin::initialize() {
        return 0;
    }

    DimsExprs CustomSlicePlugin::getOutputDimensions(
            int outputIndex, const DimsExprs *inputs, int nbInputs, IExprBuilder &exprBuilder) {
        //output the result to channel
        DimsExprs scale;
        scale.nbDims = 1;
        scale.d[0] = exprBuilder.constant(2);

        DimsExprs out_dim;
        out_dim.nbDims = 4;
        out_dim.d[0] = inputs[0].d[0];
        out_dim.d[1] = inputs[0].d[1];
        out_dim.d[2] = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *inputs[0].d[2], *scale.d[0]);
        out_dim.d[3] = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *inputs[0].d[3], *scale.d[0]);

        return out_dim;
    }

    // Set plugin namespace
    void CustomSlicePlugin::setPluginNamespace(const char *pluginNamespace) {
        mPluginNamespace = pluginNamespace;
    }

    const char *CustomSlicePlugin::getPluginNamespace() const {
        return mPluginNamespace;
    }

    // Return the DataType of the plugin output at the requested index
    DataType CustomSlicePlugin::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const {
//        return DataType::kFLOAT;
        return inputTypes[0];
    }

    void CustomSlicePlugin::configurePlugin(const DynamicPluginTensorDesc *in, int nbInputs,
                                            const DynamicPluginTensorDesc *out, int nbOutputs) {
    }

    const char *CustomSlicePlugin::getPluginType() const {
        return "CustomSliceLayer_TRT";
    }

    const char *CustomSlicePlugin::getPluginVersion() const {
        return "1";
    }

    void CustomSlicePlugin::destroy() {
        delete this;
    }

    // Clone the plugin
    IPluginV2DynamicExt *CustomSlicePlugin::clone() const {
        CustomSlicePlugin *p = new CustomSlicePlugin(mStart, mStride);
        p->setPluginNamespace(mPluginNamespace);
        return p;
    }

    __global__ void kernel_slice_persistent(
            nvinfer1::DataType data_type,
            const int *start,
            const int *size,
            const int *stride,
            const int *input_size,
            const float *input_data,
            float *output_data
    ) {
        int istart_b = start[0], istart_c = start[1], istart_h = start[2], istart_w = start[3];
        int istride_b = stride[0], istride_c = stride[1], istride_h = stride[2], istride_w = stride[3];
        int isize_b = size[0], isize_c = size[1], isize_h = size[2], isize_w = size[3];
        int input_size_b = input_size[0], input_size_c = input_size[1], input_size_h = input_size[2], input_size_w = input_size[3];

        int istart_y_idx = blockIdx.y * blockDim.y + threadIdx.y;
        int istart_x_idx = blockIdx.x * blockDim.x + threadIdx.x;
        int istep_y = gridDim.y * blockDim.y;
        int istep_x = gridDim.x * blockDim.x;
        int output_chw = isize_c * isize_h * isize_w;
        int output_hw = isize_h * isize_w;
        int output_w = isize_w;
        int input_chw = input_size_c * input_size_h * input_size_w;
        int input_hw = input_size_h * input_size_w;
        int input_w = input_size_w;
        for (int ib = 0; ib < isize_b; ib += istride_b) {
            for (int ic = 0; ic < isize_c; ic += istride_c) {
                for (int ih = istart_y_idx; ih < isize_h; ih += istep_y) {
                    for (int iw = istart_x_idx; iw < isize_w; iw += istep_x) {
                        int output_index = ib * output_chw + ic * output_hw + ih * output_w + iw;
                        int input_ib = ib * istride_b + istart_b;
                        int input_ic = ic * istride_c + istart_c;
                        int input_ih = ih * istride_h + istart_h;
                        int input_iw = iw * istride_w + istart_w;
                        int input_index = input_ib * input_chw + input_ic * input_hw + input_ih * input_w + input_iw;
                        output_data[output_index] = input_data[input_index];
                    }
                }
            }
        }

        return;
    }

    __global__ void kernel_slice(
            nvinfer1::DataType data_type,
            const int *start,
            const int *size,
            const int *stride,
            const int *input_size,
            const float *input_data,
            float *output_data
    ) {
        int istart_b = start[0], istart_c = start[1], istart_h = start[2], istart_w = start[3];
        int istride_b = stride[0], istride_c = stride[1], istride_h = stride[2], istride_w = stride[3];
        int isize_b = size[0], isize_c = size[1], isize_h = size[2], isize_w = size[3];
        int input_size_b = input_size[0], input_size_c = input_size[1], input_size_h = input_size[2], input_size_w = input_size[3];

        int istart_y_idx = blockIdx.y * blockDim.y + threadIdx.y;
        int istart_x_idx = blockIdx.x * blockDim.x + threadIdx.x;
        int output_chw = isize_c * isize_h * isize_w;
        int output_hw = isize_h * isize_w;
        int output_w = isize_w;
        int input_chw = input_size_c * input_size_h * input_size_w;
        int input_hw = input_size_h * input_size_w;
        int input_w = input_size_w;
        for (int ib = 0; ib < isize_b; ib += istride_b) {
            for (int ic = 0; ic < isize_c; ic += istride_c) {
                int ih = istart_y_idx, iw = istart_x_idx;
                if (ih < isize_h && iw < isize_w) {
                    int output_index = ib * output_chw + ic * output_hw + ih * output_w + iw;
                    int input_ib = ib * istride_b + istart_b;
                    int input_ic = ic * istride_c + istart_c;
                    int input_ih = ih * istride_h + istart_h;
                    int input_iw = iw * istride_w + istart_w;
                    int input_index = input_ib * input_chw + input_ic * input_hw + input_ih * input_w + input_iw;
                    output_data[output_index] = input_data[input_index];
                }
            }
        }

        return;
    }

    void slice_func(nvinfer1::DataType data_type,
                    const PluginTensorDesc *inputDesc,
                    const PluginTensorDesc *outputDesc,
                    const float *input_data,
                    float *output_data,
                    int *mStart,
                    int *mStride,
                    cudaStream_t stream) {
        Dims input_size = inputDesc[0].dims;
        DataType input_type = inputDesc[0].type;
        Dims size = outputDesc[0].dims;
        DataType output_type = outputDesc[0].type;
        int output_b = size.d[0], output_c = size.d[1], output_h = size.d[2], output_w = size.d[3];
        int ThreadXNumPerBlock = 32, ThreadYNumPerBlock = 32;
        int BlockXNum = output_w / ThreadXNumPerBlock + 1;
        int BlockYNum = output_h / ThreadYNumPerBlock + 1;
        dim3 blockSize(ThreadXNumPerBlock, ThreadXNumPerBlock);
        dim3 gridSize(BlockXNum, BlockYNum);

        size_t len = sizeof(int) * 4;
        int *dev_Start = nullptr;
        int *dev_Stride = nullptr;
        int *dev_Size = nullptr;
        int *dev_InputSize = nullptr;
        CUDA_CHECK(cudaMalloc((void **) &dev_Start, len));
        CUDA_CHECK(cudaMalloc((void **) &dev_Stride, len));
        CUDA_CHECK(cudaMalloc((void **) &dev_Size, len));
        CUDA_CHECK(cudaMalloc((void **) &dev_InputSize, len));
        CUDA_CHECK(cudaMemcpy(dev_Start, mStart, len, cudaMemcpyHostToDevice));
        CUDA_CHECK(cudaMemcpy(dev_Stride, mStride, len, cudaMemcpyHostToDevice));
        CUDA_CHECK(cudaMemcpy(dev_Size, size.d, len, cudaMemcpyHostToDevice));
        CUDA_CHECK(cudaMemcpy(dev_InputSize, input_size.d, len, cudaMemcpyHostToDevice));

        kernel_slice<<<gridSize, blockSize, 0, stream>>>(data_type, dev_Start, dev_Size, dev_Stride, dev_InputSize,
                                                                   input_data, output_data);

        CUDA_CHECK(cudaPeekAtLastError());
        CUDA_CHECK(cudaFree(dev_Start));
        CUDA_CHECK(cudaFree(dev_Stride));
        CUDA_CHECK(cudaFree(dev_Size));
        CUDA_CHECK(cudaFree(dev_InputSize));
    }

    // TODO: 还不能支持半精度计FP16计算！！！
    int CustomSlicePlugin::enqueue(
            const PluginTensorDesc *inputDesc,
            const PluginTensorDesc *outputDesc,
            const void *const *inputs,
            void *const *outputs,
            void *workspace,
            cudaStream_t stream) {
        auto data_type = nvinfer1::DataType::kFLOAT;
        slice_func(data_type, inputDesc, outputDesc, (const float *)inputs[0], (float *)outputs[0],
                mStart, mStride, stream);

        return 0;
    }

    bool CustomSlicePlugin::supportsFormatCombination(
            int pos, const PluginTensorDesc *inOut, int nbInputs, int nbOutputs) {
        return true;
    }

    size_t CustomSlicePlugin::getWorkspaceSize(
            const PluginTensorDesc *inputs, int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const {
        Dims input_size = inputs[0].dims;
        Dims output_size = outputs[0].dims;
        auto data_type = inputs[0].type;

        size_t worksize = 0;
        switch (data_type) {
            case nvinfer1::DataType::kFLOAT:
                worksize += input_size.d[0] * input_size.d[1] * input_size.d[2] * input_size.d[3] * sizeof(float);
                worksize += output_size.d[0] * output_size.d[1] * output_size.d[2] * output_size.d[3] * sizeof(float);
                break;
            case nvinfer1::DataType::kHALF:
                worksize += input_size.d[0] * input_size.d[1] * input_size.d[2] * input_size.d[3] * sizeof(float) / 2;
                worksize +=
                        output_size.d[0] * output_size.d[1] * output_size.d[2] * output_size.d[3] * sizeof(float) / 2;
                break;
            default:
                worksize += 1;
                break;
        }

        worksize += sizeof(int) * 4 * 4;

        return worksize;
    }


    //////////////////////////////////////////////////////////////////////////////////
    PluginFieldCollection CustomSlicePluginCreator::mFC{};
    std::vector <PluginField> CustomSlicePluginCreator::mPluginAttributes;

    CustomSlicePluginCreator::CustomSlicePluginCreator() {
        mPluginAttributes.clear();

        mFC.nbFields = mPluginAttributes.size();
        mFC.fields = mPluginAttributes.data();
    }

    const char *CustomSlicePluginCreator::getPluginName() const {
        return "CustomSliceLayer_TRT";
    }

    const char *CustomSlicePluginCreator::getPluginVersion() const {
        return "1";
    }

    const PluginFieldCollection *CustomSlicePluginCreator::getFieldNames() {
        return &mFC;
    }

    IPluginV2DynamicExt *CustomSlicePluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) {
        assert(fc->nbFields == 2);
        assert(strcmp(fc->fields[0].name, "start") == 0);
        assert(strcmp(fc->fields[1].name, "stride") == 0);
        int *start = (int *) (fc->fields[0].data);
        int *stride = (int *) (fc->fields[1].data);

        CustomSlicePlugin *obj = new CustomSlicePlugin(start, stride);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }

    IPluginV2DynamicExt *
    CustomSlicePluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) {
        // This object will be deleted when the network is destroyed, which will
        // call CustomSlicePlugin::destroy()
        CustomSlicePlugin *obj = new CustomSlicePlugin(serialData, serialLength);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }


}

